import cv2
import mss
import numpy as np
import pyautogui
from pynput.keyboard import Controller as KC
from pynput.keyboard import Key
pyautogui.PAUSE = 0


class Action:
    def __init__(self):
        self.kc = KC()

    def act(self, action):
        if action == 0: self.kc.tap(Key.shift)
        elif action == 1: self.kc.tap("a")
        elif action == 2: self.kc.tap("s")
        elif action == 3: self.kc.tap("d")

    def temp_act(self, action):
        if action[0] == 0:  # press
            self.kc.press(action[1])
        elif action[0] == 1:  # release
            self.kc.release(action[1])
        elif action[0] == 2:  # Button.left
            pyautogui.click(action[1][0], action[1][1])
        elif action[0] == 3:  # Button.right
            pyautogui.rightClick(action[1][0], action[1][1])


class Screen:
    def __init__(self):
        self.cap = mss.mss()
        left, top, width, height = pyautogui.getActiveWindow().box
        left, top, right, bottom = left+9, top, left+width-9, top+height-9
        self.region = (left, top, right, bottom)
        self.region = (left, top+38, right, bottom)  # 原神标题裁剪

    def shot(self):
        img = self.cap.grab(self.region)
        img = np.resize(np.array(img.raw), (img.height, img.width, 4))  # BGRA四通道
        img = cv2.resize(img, (96, 96))
        img = cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)
        return [np.dot(img[..., :], [0.299, 0.587, 0.114]) / 127]  # 前期研究简化：注意将灰度通道体现出来
        # return img // 5  # 参照人类标准，控制色彩细粒度


class Env:
    def __init__(self):
        self.screen, self.action = Screen(), Action()

    def reset(self):
        return self.screen.shot()

    def step(self, action):  # TODO 暂时配合第一阶段测试，鼓励模型朝一个方向移动
        self.action.act(action)
        if action == 1: return self.screen.shot(), 100
        return self.screen.shot(), 0


def show(img):
    cv2.namedWindow("test")  # 创建一个图片大小（cv2.WINDOW_NORMAL）的窗口
    cv2.imshow("test", img)  # 在该窗口显示图片
    cv2.waitKey(0)  # 等待Key关闭窗口
    cv2.destroyAllWindows()


if __name__ == "__main__":
    vision = Screen()
    img = vision.shot() * 5
    show(img)

